package com.zon.len.dynamic;

import com.alibaba.druid.pool.DruidDataSource;
import com.zon.len.dynamic.exception.DynamicDataSourceException;
import com.zon.len.dynamic.properties.MultipleDruidDataSourceProperties;
import com.zon.len.dynamic.properties.MultipleDruidDataSourceProperties.DynamicDruidDataSourceProperties;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

/**
 * @author ZonLen since on 2021/12/31 上午10:26
 */
@Slf4j
public class DynamicDruidDataSource extends AbstractRoutingDataSource implements
    ApplicationContextAware, DisposableBean {

  private DruidDataSourceCreator druidDataSourceCreator;

  private final Map<DynamicDataSourceKey, DruidDataSource> druidDataSources = new HashMap<>(16);

  private final MultipleDruidDataSourceProperties multipleDruidDataSourceProperties;

  /**
   * 资源销毁
   */
  @Override
  public void destroy() {
    druidDataSources.forEach((key, dataSource) -> {
      dataSource.close();
      log.info("Dynamic DataSource {} is destroyed", key.name());
    });
  }

  public DynamicDruidDataSource(
      MultipleDruidDataSourceProperties multipleDruidDataSourceProperties,
      DataSourceProperties dataSourceProperties) {
    //if not found prefix 'spring.datasource.druid' jdbc properties ,'spring.datasource' prefix jdbc properties will be used.
    if (multipleDruidDataSourceProperties.getUsername() == null) {
      multipleDruidDataSourceProperties.setUsername(dataSourceProperties.determineUsername());
    }
    if (multipleDruidDataSourceProperties.getPassword() == null) {
      multipleDruidDataSourceProperties.setPassword(dataSourceProperties.determinePassword());
    }
    if (multipleDruidDataSourceProperties.getUrl() == null) {
      multipleDruidDataSourceProperties.setUrl(dataSourceProperties.determineUrl());
    }
    if (multipleDruidDataSourceProperties.getDriverClassName() == null) {
      multipleDruidDataSourceProperties
          .setDriverClassName(dataSourceProperties.getDriverClassName());
    }
    this.multipleDruidDataSourceProperties = multipleDruidDataSourceProperties;

  }

  /**
   * 初始化操作
   */
  @Override
  public void afterPropertiesSet() {
    loadDefaultDataSource();
    loadDynamicDataSources();
    log.info("Dynamic datasource [" + druidDataSources.keySet().stream()
        .map(DynamicDataSourceKey::name).collect(
            Collectors.joining(",")) + "] initialized");
    Map<Object, Object> targetDataSources = new HashMap<>(druidDataSources.size());
    druidDataSources.forEach(targetDataSources::put);
    setTargetDataSources(targetDataSources);
    super.afterPropertiesSet();
  }

  private void loadDefaultDataSource() {
    final DruidDataSource defaultDataSource = druidDataSourceCreator
        .createDataSource(multipleDruidDataSourceProperties);
    druidDataSources.put(DynamicDataSourceKey.DEFAULT_DB, defaultDataSource);
    setDefaultTargetDataSource(defaultDataSource);
  }

  private void loadDynamicDataSources() {
    final List<DynamicDruidDataSourceProperties> dynamicDataSourcePropertiesList = multipleDruidDataSourceProperties
        .getDynamicSources();
    if (CollectionUtils.isNotEmpty(dynamicDataSourcePropertiesList)) {
      for (DynamicDruidDataSourceProperties dynamicDruidDataSourceProperties : dynamicDataSourcePropertiesList) {
        final String sourceKey = dynamicDruidDataSourceProperties.getSourceKey();
        druidDataSources.put(DynamicDataSourceKey.sourceKey(sourceKey),
            druidDataSourceCreator.createDataSource(dynamicDruidDataSourceProperties));
      }
    }
  }

  @Override
  public void setTargetDataSources(Map<Object, Object> targetDataSources) {
    super.setTargetDataSources(targetDataSources);
    super.afterPropertiesSet();// 必须添加该句，否则新添加数据源无法识别到
  }

  @Override
  protected Object determineCurrentLookupKey() {
    return DataSourceContextHolder.getDataSourceKey();
  }

  @Override
  public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
    this.druidDataSourceCreator = new DruidDataSourceCreator(applicationContext);
  }


  public enum DynamicDataSourceKey {
    /**
     * 默认的DB方式
     */
    DEFAULT_DB,
    DEMO_DB1,
    BILL_CENTER_DB;

    public static DynamicDataSourceKey sourceKey(String sourceKey) {
      for (DynamicDataSourceKey dynamicDataSourceKey : DynamicDataSourceKey.values()) {
        if (dynamicDataSourceKey.name().equals(sourceKey)) {
          return dynamicDataSourceKey;
        }
      }
      throw new DynamicDataSourceException(
          "Enum DynamicDataSourceKey not find sourceKey:" + sourceKey);
    }
  }
}
